// $Id: wfc.c 480 2014-04-30 22:06:15Z alan.watson.f $

#define _POSIX_C_SOURCE 200112L

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/tcp.h>
#include <sys/time.h>

#include "oaxaca-config.h"
#include "oaxaca.h"

#include "minmax.h"
#include "mprintf.h"

////////////////////////////////////////////////////////////////////////////////

int
oaxaca_wfc_initialized(void)
{
  return oaxaca_shm != NULL && oaxaca_shm->wfc.initialized;
}

char
oaxaca_wfc_command(void)
{
  if (oaxaca_shm == NULL)
    return 0;
  else
    return oaxaca_shm->wfc.command;
}

void
oaxaca_wfc_set_command(char command)
{
  if (oaxaca_shm != NULL)
    oaxaca_shm->wfc.command = command;
}

double
oaxaca_wfc_a_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.a[i];
  else
    return 0.0;
}

void
oaxaca_wfc_a_flat_vector_set(unsigned int i, double a)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    oaxaca_shm->wfc.aflat[i] = a;
}

double
oaxaca_wfc_a_flat_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.aflat[i];
  else
    return 0.0;
}

void
oaxaca_wfc_a_vector_set(unsigned int i, double a)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    oaxaca_shm->wfc.a[i] = a;
}

double
oaxaca_wfc_a_mean_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.amean[i];
  else
    return 0.0;
}

double
oaxaca_wfc_a_sigma_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.asigma[i];
  else
    return 0.0;
}

double
oaxaca_wfc_da_mean_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.damean[i];
  else
    return 0.0;
}

double
oaxaca_wfc_da_sigma_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_NA_MAX)
    return oaxaca_shm->wfc.dasigma[i];
  else
    return 0.0;
}

double
oaxaca_wfc_tm_v_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_TM_NV)
    return oaxaca_shm->wfc.tmv[i];
  else
    return 0.0;
}

double
oaxaca_wfc_dm_v_vector_ref(unsigned int i)
{
  if (oaxaca_shm != NULL && i < OAXACA_DM_NV)
    return oaxaca_shm->wfc.dmv[i];
  else
    return 0.0;
}

////////////////////////////////////////////////////////////////////////////////

#if USE_TCP

static FILE *out_tm;
static FILE *in_tm;
static FILE *out_dm;
static FILE *in_dm;

static void
server_open(const char *servername, int serverport, FILE **in, FILE **out)
{
  oaxaca_info("connecting to %s:%d.", servername, serverport);

  int fd = socket(AF_INET, SOCK_STREAM, 0);
  if (fd < 0)
    oaxaca_error("socket failed.");
    
  struct sockaddr_in sa;
  memset(&sa, 0, sizeof(sa));
  sa.sin_family = AF_INET;
  sa.sin_port = htons(serverport);
  
  errno = 0;
  struct hostent *hostinfo = gethostbyname(servername);
  if (hostinfo == NULL)
    oaxaca_error("gethostbyname failed: %s.", strerror(errno));
  sa.sin_addr = * (struct in_addr *) hostinfo->h_addr_list[0];
    
  errno = 0;
  if (connect(fd, (struct sockaddr *) &sa, sizeof(sa)) < 0)
    oaxaca_error("connect failed: %s.", strerror(errno));

  int flag = 1;
  int result = setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (char *) &flag, sizeof(int));
  if (result < 0)
    oaxaca_error("setsockopt failed: %s.", strerror(errno));

  *in = fdopen(fd, "r");
  *out = fdopen(fd, "w");
  setvbuf(*out, NULL, _IOLBF, 0);

  oaxaca_info("connected.", servername, serverport);
  
}

#else

static int fd_tm;
static struct sockaddr_in sa_tm;
static int fd_dm;
static struct sockaddr_in sa_dm;

static void
server_open(const char *servername, int serverport, int *fd, struct sockaddr_in *sa)
{
  errno = 0;
  *fd = socket(AF_INET, SOCK_DGRAM, 0);
  if (*fd < 0) {
    oaxaca_error("socket failed: %s.", strerror(errno));
  }
  
  memset(sa, 0, sizeof(*sa));
  sa->sin_family = AF_INET;
  sa->sin_port = htons(serverport);
  errno = 0;
  struct hostent *hostinfo = gethostbyname(servername);
  if (hostinfo == NULL) {
    oaxaca_error("gethostbyname failed: %s.", strerror(errno));
  }
  sa->sin_addr = * (struct in_addr *) hostinfo->h_addr_list[0];
  
}

#endif

////////////////////////////////////////////////////////////////////////////////


static char *
format_m_line(int n, unsigned long u[n])
{
  char *curr = mprintf("%s", "m");
  char *prev;
  for (int i = 0; i < n; ++i) {
    prev = curr;
    curr = mprintf("%s %04lx", prev, u[i]);
    free(prev);
  }
  prev = curr;
  curr = mprintf("%s\n", prev);
  free(prev);
  return curr;  
}

__attribute__((unused))
static void
wait_dm(void)
{
#if USE_TCP
  unsigned long reply;
  if (fscanf(in_dm, "%lu", &reply) != 1)
    oaxaca_error("invalid reply from dm.");
#endif
}

static void
set_dm_voltages(void)
{
  double z[64];
  for (int i = 0; i < 64; ++i)
    z[i] = 0.0;
  
  double pi = 4.0 * atan(1.0);

  // The tilt electrodes in order with origin 0. From the "finger" counter clockwise.
  int fr_tilt_electrode[] = {
    39, 37, 35, 33, 34, 36, 38, 7, 5, 3, 1, 0, 2, 4, 6
  };
  int bk_tilt_electrode[] = {
    57, 61, 62, 58, 54, 50, 44, 40, 11, 15, 19, 23, 27, 31, 26
  };

  double fr_tilt_theta[15];
  double bk_tilt_theta[15];
  for (int i = 0; i < 15; ++i) {
    fr_tilt_theta[i] = (i + 0.50) * 2.0 * pi / 15.0;
    bk_tilt_theta[i] = (i + 0.25) * 2.0 * pi / 15.0;
  }

#if 0
  // Mode 0 is x-tilt
  {
    double a = oaxaca_shm->wfc.a[0];
    for (int i = 0; i < 15; ++i) {
      double w = cos(1.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = cos(1.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }
    
  // Mode 1 is y-tilt
  {
    double a = oaxaca_shm->wfc.a[1];
    for (int i = 0; i < 15; ++i) {
      double w = sin(1.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = sin(1.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }
#endif

  // Mode 2 is Z4 defocus.
  {
    double a = oaxaca_shm->wfc.a[2];
    if (a > 0.0) {
      z[8] = +a;
      z[25] = z[24] = z[60] = z[48] = z[42] = z[17] = z[30] = z[59] = z[63] = z[56] = z[52] = z[46] = z[9] = z[13] = z[21] = z[29] = z[28] = 0.0;
    } else {
      z[8] = 0.0;
      z[25] = z[24] = z[60] = z[48] = z[42] = z[17] = z[30] = z[59] = z[63] = z[56] = z[52] = z[46] = z[9] = z[13] = z[21] = z[29] = z[28] = -a;      
    }
  }
  
  // Mode 3 is Z5 astigmatism.
  {
    double a = oaxaca_shm->wfc.a[3];
    for (int i = 0; i < 15; ++i) {
      double w = sin(2.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = sin(2.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }

  // Mode 4 is Z6 astigmatism.
  {
    double a = oaxaca_shm->wfc.a[4];
    for (int i = 0; i < 15; ++i) {
      double w = cos(2.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = cos(2.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }
  
  // Mode 5 is Z9 trefoil
  {
    double a = oaxaca_shm->wfc.a[5];
    for (int i = 0; i < 15; ++i) {
      double w = sin(3.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = sin(3.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }
  
  // Mode 6 is Z10 trefoil
  {
    double a = oaxaca_shm->wfc.a[6];
    for (int i = 0; i < 15; ++i) {
      double w = cos(3.0 * fr_tilt_theta[i]);
      if (w * a >= 0)
        z[fr_tilt_electrode[i]] += a * w;
    }
    for (int i = 0; i < 15; ++i) {
      double w = cos(3.0 * bk_tilt_theta[i]);
      if (w * a <= 0)
        z[bk_tilt_electrode[i]] -= a * w;
    }
  }
  
  double v[64];
  for (int i = 0; i < 64; ++i) {
    if (z[i] <= 0.0)
      v[i] = +0.0;
    else if (z[i] <= 1.0)
      v[i] = sqrt(z[i]);
    else
      v[i] = +1.0;
    oaxaca_shm->wfc.dmv[i] = v[i];
  }
  
  unsigned long u[64];
  for (int i = 0; i < 64; ++i) {
    u[i] = v[i] * 0x3fff;
  }

  char *line = format_m_line(64, u);

#if USE_TCP
  fputs(line, out_dm);
#else
  errno = 0;
  int r = sendto(fd_dm, line, strlen(line) + 1, 0, (struct sockaddr *)&sa_dm, sizeof(sa_dm));
  if (r < 0) {
    oaxaca_error("sendto failed: %s.", strerror(errno));
  }
#endif

  free(line);
}

__attribute__((unused))
static void
wait_tm(void)
{
#if USE_TCP
  unsigned long reply;
  if (fscanf(in_tm, "%lu", &reply) != 1)
    oaxaca_error("invalid reply from tm.");
#endif
}

static void
set_tm_voltages(void)
{
  double v[3];
  for (int i = 0; i < 3; ++i)
    v[i] = 0.0;

  double a0 = oaxaca_shm->wfc.a[0];
  double a1 = oaxaca_shm->wfc.a[1];

  // Make a1 and a2 orthogonal to a3.
//  a0 += +0.045 * oaxaca_shm->wfc.a[1];
//  a0 += -0.250 * oaxaca_shm->wfc.a[2];

#if 0
  // Three actuators.
  double pi = 4.0 * atan(1.0);
  double x[3] = { cos(0.0 * pi / 3.0), cos(2.0 * pi / 3.0), cos (4.0 * pi / 3) };
  double y[3] = { sin(0.0 * pi / 3.0), sin(2.0 * pi / 3.0), sin (4.0 * pi / 3) };
  for (int i = 0; i < 3; ++i)
    v[i] = 0.5 * x[i]* (a0 + 1.0) + 0.5 * y[i] * (a1 + 1.0);
#else
  // Only actuators 0 and 2.
  v[0] = 0.5 + 0.5 * a0 + 0.25 * a1;
  v[1] = 0.0;
  v[2] = 0.5 + 0.5 * a1;
#endif

  for (int i = 0; i < 3; ++i) {    
    if (v[i] <= 0.0)
      v[i] = +0.0;
    else if (v[i] > +1.0)
      v[i] = +1.0;
  }
  
  for (;;) {
  
    // Change the voltages by at most 0.5 V per ms. A change of 0.5 V
    // corresponds to 5% of the full range is (0,+10) V or 0.1 in the
    // normalized range of a of (-1,+1).
    
    double dvmax = 0.1;

    for (int i = 0; i < 3; ++i) {
      double oldv = oaxaca_shm->wfc.tmv[i];
      double dv = v[i] - oldv;
      dv = clip(dv, -dvmax, +dvmax);
      double newv = oldv + dv;
      oaxaca_shm->wfc.tmv[i] = newv;
    }
  
    unsigned long u[3];
    for (int i = 0; i < 3; ++i) {      
      u[i] = oaxaca_shm->wfc.tmv[i] * 0xffff;
      if (u[i] > 0xffff)
        u[i] = 0xffff;
    }
    
    char *line = format_m_line(3, u);

#if USE_TCP
    fputs(line, out_tm);
#else
    errno = 0;
    int r = sendto(fd_tm, line, strlen(line), 0, (struct sockaddr *)&sa_tm, sizeof(sa_tm));
    if (r < 0) {
      oaxaca_error("sendto failed: %s.", strerror(errno));
    }
#endif

    free(line);
      
    //wait_tm();

    int completed = 1;
    for (int i = 0; i < 3; ++i)
      if (v[i] != oaxaca_shm->wfc.tmv[i])
        completed = 0;
        
    if (completed)
      break;
      
    oaxaca_info("rate limiting TM.");

    oaxaca_usleep(1000);
      
  }

}

static void
set_voltages(void)
{
  set_dm_voltages();
  set_tm_voltages();
}

////////////////////////////////////////////////////////////////////////////////

static void
atexitcallee(void)
{
  if (oaxaca_shm != NULL)
    oaxaca_shm->wfc.initialized = 0;
}

void
oaxaca_wfc_main(void)
{
  oaxaca_common_main("wfc");
  atexit(atexitcallee);
  
  oaxaca_info("connecting to wfc hardware.");
#if USE_TCP
  server_open("dm", 5000, &in_dm, &out_dm);
  server_open("tm", 5001, &in_tm, &out_tm);
#else
  server_open("dm", 5000, &fd_dm, &sa_dm);
  server_open("tm", 5001, &fd_tm, &sa_tm);
#endif

  oaxaca_shm->wfc.initialized = 0;
  oaxaca_shm->wfc.command     = 't';
  oaxaca_shm->wfc.totaldelay  = 0.0;
  for (int i = 0; i < OAXACA_TM_NV; ++i)
    oaxaca_shm->wfc.tmv[i] = 0.0;
  for (int i = 0; i < OAXACA_DM_NV; ++i)
    oaxaca_shm->wfc.dmv[i] = 0.0;
  for (int i = 0; i < OAXACA_NA_MAX; ++i) {
    oaxaca_shm->wfc.a[i]          = 0.0;
    oaxaca_shm->wfc.aflat[i]      = 0.0;
    oaxaca_shm->wfc.amean[i]      = 0.0;
    oaxaca_shm->wfc.asigma[i]     = 0.0;
    oaxaca_shm->wfc.damean[i]     = 0.0;
    oaxaca_shm->wfc.dasigma[i]    = 0.0;
    oaxaca_shm->wfc.asaturated[i] = 0.0;
  }
  oaxaca_shm->wfc.initialized = 1;
    
  oaxaca_info("starting command loop.");

  for (;;) {
    int command = oaxaca_shm->wfc.command;
    switch (command) {
    case 'i': // start idling
      oaxaca_shm->wfc.command = 'I';
      break;
    case 'I': // continue idling
      break;
    case 'f': // fixed
      for (int i = 0; i < OAXACA_NA_MAX; ++i) {
        oaxaca_shm->wfc.amean[i]  = oaxaca_shm->wfc.a[i];
        oaxaca_shm->wfc.asigma[i] = 0.0;
      }
      set_voltages();        
      oaxaca_shm->wfc.command = 'I';
      break;      
    case 't': // trigger
      {
        set_voltages();
        struct timeval tv;
        gettimeofday(&tv, NULL);
        double timestamp = tv.tv_sec + tv.tv_usec * 1e-6;
        oaxaca_shm->wfc.totaldelay += timestamp - oaxaca_shm->wfc.timestamp;
        oaxaca_shm->wfc.command = 'I';
      }
      break;
    default:
      oaxaca_error("unrecognized command \'%c\'.", command);
      break;
    }
    oaxaca_usleep(200);
  }
  
  oaxaca_error("command loop terminated.");
}

////////////////////////////////////////////////////////////////////////////////